﻿
#define SPADAS_DEPRECATED_HIDE
#include "NextFileIO.h"
#include "NextUtility.h"
#include "ffmpegdl.h"

namespace videonext
{
	enum class VideoContainerType
	{
		Invalid = 0,
		AVI = 1,
		MP4 = 2,
		MKV = 3,
	};

	class FileReaderVars : public Vars
	{
	public:
		Bool peekMode;
		Array<Char> utf8PathArray;
		FileFormat format;
		VideoContainerType containerType;
		Pointer formatContext;
		Int targetStreamIndex;
		Bool mkvZeroDts;
	};

	class FileWriterVars : public Vars
	{
	public:
		FileFormat format;
		VideoContainerType containerType;
		Pointer formatContext;
		UInt count;
		Binary lastFrame;
		Bool writeHeaderLater;
	};

	Array<Char> pathToUtf8CharArray(Path path)
	{
		Binary utf8Binary = path.fullPath().toBinary();
		Array<Char> chars(utf8Binary.size() + 1, 0);
		utility::memoryCopy(utf8Binary.data(), chars.data(), utf8Binary.size());
		return chars;
	}

	Int getH264ExtradataSize(Binary frame)
	{
		if (frame.size() < 4) return -1;

		Bool found67 = FALSE, found68 = FALSE;
		UInt searchRange = frame.size() - 4;
		for (UInt i = 0; i < searchRange; i++)
		{
			if (frame[i] == 0 && frame[i + 1] == 0 && frame[i + 2] == 1)
			{
				if (frame[i + 3] == 0x67 || frame[i + 3] == 0x27) found67 = TRUE;
				else if (frame[i + 3] == 0x68 || frame[i + 3] == 0x28) found68 = TRUE;
				else if (found67 && found68) return i;
			}
		}
		return -1;
	}

	Int getH265ExtradataSize(Binary frame)
	{
		if (frame.size() < 4) return -1;

		Bool found40 = FALSE, found42 = FALSE, found44 = FALSE;
		UInt searchRange = frame.size() - 4;
		for (UInt i = 0; i < searchRange; i++)
		{
			if (frame[i] == 0 && frame[i + 1] == 0 && frame[i + 2] == 1)
			{
				if (frame[i + 3] == 0x40) found40 = TRUE;
				else if (frame[i + 3] == 0x42) found42 = TRUE;
				else if (frame[i + 3] == 0x44) found44 = TRUE;
				else if (found40 && found42 && found44) return i;
			}
		}
		return -1;
	}
}

using namespace videonext;

FileReader::FileReader() : Object<FileReaderVars>(new FileReaderVars, TRUE)
{
	vars->peekMode = FALSE;
	vars->formatContext = NULL;
}
Bool FileReader::open(Path filePath, FileFormat& format)
{
	close();

	if (filePath.isFolder() || !filePath.exist()) return FALSE;

	String ext = filePath.extension().toLower();
	if (ext == ".avi") vars->containerType = VideoContainerType::AVI;
	else if (ext == ".mp4") vars->containerType = VideoContainerType::MP4;
	else if (ext == ".mkv") vars->containerType = VideoContainerType::MKV;
	else return FALSE;

	vars->utf8PathArray = pathToUtf8CharArray(filePath);

	Pointer formatContext = NULL;
	Int resultOpenInput = avformat_open_input(&formatContext, vars->utf8PathArray.data(), NULL, NULL);
	if (resultOpenInput < 0) return FALSE;

	if (!vars->peekMode)
	{
		Int resultFindStreamInfo = avformat_find_stream_info(formatContext, NULL);
		if (resultFindStreamInfo < 0)
		{
			avformat_free_context(formatContext);
			return FALSE;
		}
	}

	Int targetStreamIndex = -1;
	for (UInt i = 0; i < avformat_get_nb_streams(formatContext); i++)
	{
		Pointer stream = avformat_get_streams(formatContext)[i];
		if (getFFMPEGLibraryVersion() >= 3)
		{
			if (avformat_stream_get_codec_type(stream) != AVMediaType::AVMEDIA_TYPE_VIDEO) continue;
		}
		else
		{
			Pointer codecCtx = avformat_stream_get_codec(stream);
			if (avcodec_get_codec_type(codecCtx) != AVMediaType::AVMEDIA_TYPE_VIDEO) continue;
		}
		targetStreamIndex = i;
		break;
	}
	if (targetStreamIndex == -1)
	{
		avformat_free_context(formatContext);
		return FALSE;
	}

	Pointer targetStream = avformat_get_streams(formatContext)[targetStreamIndex];
	if (getFFMPEGLibraryVersion() >= 3)
	{
		if (avformat_stream_get_codec_id(targetStream) == AVCodecID::AV_CODEC_ID_H264) format.codec = VideoDataCodec::H264;
		else if (avformat_stream_get_codec_id(targetStream) == AVCodecID::AV_CODEC_ID_MJPEG) format.codec = VideoDataCodec::MJPEG;
		else if (avformat_stream_get_codec_id(targetStream) == AVCodecID::AV_CODEC_ID_Y41P) format.codec = VideoDataCodec::YUV411;
		else if (avformat_stream_get_codec_id(targetStream) == AVCodecID::AV_CODEC_ID_YUV4) format.codec = VideoDataCodec::YUV420;
		else if (avformat_stream_get_codec_id(targetStream) == AVCodecID::AV_CODEC_ID_H265) format.codec = VideoDataCodec::H265;
		else format.codec = VideoDataCodec::Invalid;
		format.width = avformat_stream_get_width(targetStream);
		format.height = avformat_stream_get_height(targetStream);
		Int num = 0, den = 0;
		avformat_stream_get_time_base(targetStream, num, den);
		format.fps = den / num;
		if (vars->containerType == VideoContainerType::MP4) format.fps /= 1000;
	}
	else
	{
		Pointer codecCtx = avformat_stream_get_codec(targetStream);
		if (avcodec_get_codec_id(codecCtx) == AVCodecID::AV_CODEC_ID_H264) format.codec = VideoDataCodec::H264;
		else if (avcodec_get_codec_id(codecCtx) == AVCodecID::AV_CODEC_ID_MJPEG) format.codec = VideoDataCodec::MJPEG;
		else if (avcodec_get_codec_id(codecCtx) == AVCodecID::AV_CODEC_ID_Y41P) format.codec = VideoDataCodec::YUV411;
		else if (avcodec_get_codec_id(codecCtx) == AVCodecID::AV_CODEC_ID_YUV4) format.codec = VideoDataCodec::YUV420;
		else if (avcodec_get_codec_id(codecCtx) == AVCodecID::AV_CODEC_ID_H265) format.codec = VideoDataCodec::H265;
		else format.codec = VideoDataCodec::Invalid;
		format.width = avcodec_get_width(codecCtx);
		format.height = avcodec_get_height(codecCtx);
		Int num = 0, den = 0;
		avformat_stream_get_time_base(targetStream, num, den);
		format.fps = den / num;
		if (vars->containerType == VideoContainerType::MP4) format.fps /= 1000;
	}

	vars->format = format;
	vars->formatContext = formatContext;
	vars->targetStreamIndex = targetStreamIndex;

	vars->mkvZeroDts = FALSE;
	if (vars->containerType == VideoContainerType::MKV) // load some frames to determine fps
	{
		vars->format.fps = 0; // make readFrame output raw dts

		UInt readFrameCount = 0;
		Binary dummy;
		UInt dts = 0;
		while (TRUE)
		{
			if (!readFrame(dummy, &dts)) break;
			if (dts == 0) vars->mkvZeroDts = TRUE;
			readFrameCount++;
			if (readFrameCount >= 10) break;
		}
		if (readFrameCount < 2 || dts == 0)
		{
			avformat_free_context(formatContext);
			vars->formatContext = NULL;
			return FALSE;
		}

		vars->format.fps = format.fps = math::round(1000.0f / ((Float)dts / (readFrameCount - (vars->mkvZeroDts ? 1 : 0))));

		// reopen
		avformat_close_input(&formatContext);
		formatContext = NULL;
		avformat_open_input(&formatContext, vars->utf8PathArray.data(), NULL, NULL);
		avformat_find_stream_info(formatContext, NULL);
		vars->formatContext = formatContext;
	}

	return TRUE;
}
Bool FileReader::seek(UInt frameIndex, Bool any)
{
	if (vars->formatContext == NULL) return FALSE;

	if (vars->containerType == VideoContainerType::AVI || vars->containerType == VideoContainerType::MP4) // ok to seek directly
	{
		Long target = (Long)frameIndex;
		if (vars->containerType == VideoContainerType::MP4) target *= 1000;

		Int resultSeek = av_seek_frame(vars->formatContext, vars->targetStreamIndex, target, any || frameIndex == 0 ? AVSEEK_FLAG_ANY : AVSEEK_FLAG_BACKWARD);
		if (resultSeek < 0) return FALSE;

		return TRUE;
	}
	else if (vars->containerType == VideoContainerType::MKV)
	{
		if (vars->mkvZeroDts) return FALSE;
		else frameIndex++;

		if (vars->format.codec == VideoDataCodec::H264 || vars->format.codec == VideoDataCodec::H265) // ok to seek directly
		{
			Long target = (Long)frameIndex * 1000 / vars->format.fps;

			Int resultSeek = av_seek_frame(vars->formatContext, vars->targetStreamIndex, target, any || frameIndex == 0 ? AVSEEK_FLAG_ANY : AVSEEK_FLAG_BACKWARD);
			if (resultSeek < 0) return FALSE;

			return TRUE;
		}
		else if (vars->format.codec == VideoDataCodec::MJPEG || vars->format.codec == VideoDataCodec::YUV411 || vars->format.codec == VideoDataCodec::YUV420)
		{
			// try to seek
			Int curStep = 10;
			Int targetIndex = (Int)frameIndex - curStep;
			while (targetIndex > 10)
			{
				Long target = (Long)targetIndex * 1000 / vars->format.fps;
				av_seek_frame(vars->formatContext, vars->targetStreamIndex, target, AVSEEK_FLAG_ANY);

				Binary dummy;
				UInt curIndex = UINF;
				if (readFrame(dummy, &curIndex) && curIndex < frameIndex)
				{
					for (UInt i = 0; i < frameIndex - curIndex - 1; i++)
					{
						readFrame(dummy);
					}
					return TRUE;
				}

				curStep *= 2;
				targetIndex = (Int)frameIndex - curStep;
			}

			// reopen and read to target
			avformat_close_input(&vars->formatContext);
			Pointer formatContext = NULL;
			avformat_open_input(&formatContext, vars->utf8PathArray.data(), NULL, NULL);
			avformat_find_stream_info(formatContext, NULL);
			vars->formatContext = formatContext;

			Binary dummy;
			for (UInt i = 0; i < frameIndex; i++)
			{
				readFrame(dummy);
			}

			return TRUE;
		}
	}

	return FALSE;
}
Bool FileReader::readFrame(Binary& frame, UInt *frameIndex)
{
	if (vars->formatContext == NULL) return FALSE;

	Bool frameLoaded = FALSE;
	while (!frameLoaded)
	{
		Pointer packet = NULL;
		if (getFFMPEGLibraryVersion() >= 3)
		{
			packet = av_packet_alloc();
		}
		else
		{
			packet = av_mallocz_av_packet();
			av_init_packet(packet);
		}

		Int resultReadFrame = av_read_frame(vars->formatContext, packet);
		if (resultReadFrame < 0) break;

		if (av_packet_get_stream_index(packet) == vars->targetStreamIndex)
		{
			Binary bin(av_packet_get_data(packet), av_packet_get_size(packet));
			frame = bin;
			if (frameIndex)
			{
				UInt index = (UInt)av_packet_get_dts(packet);
				if (vars->containerType == VideoContainerType::MP4) index /= 1000;
				else if (vars->format.fps > 0 && vars->containerType == VideoContainerType::MKV)
				{
					index = math::round((Float)index * vars->format.fps / 1000);
					if (!vars->mkvZeroDts && index > 0) index--;
				}
				*frameIndex = index;
			}
			if ((vars->containerType == VideoContainerType::MP4 || vars->containerType == VideoContainerType::MKV) && frame.size() > 4)
			{
				Bool isNormal = FALSE;
				if (frame[0] == 0 && frame[1] == 0 && frame[2] == 0 && frame[3] == 1) isNormal = TRUE;
				if (!isNormal)
				{
					UInt curIndex = 0;
					while (frame.size() > curIndex + 4)
					{
						UInt segmentSize = ((UInt)frame[curIndex] << 24) | ((UInt)frame[curIndex + 1] << 16) | ((UInt)frame[curIndex + 2] << 8) | (UInt)frame[curIndex + 3];
						frame[curIndex] = 0;
						frame[curIndex + 1] = 0;
						frame[curIndex + 2] = 0;
						frame[curIndex + 3] = 1;
						curIndex += 4 + segmentSize;
					}
				}
			}
			frameLoaded = TRUE;
		}

		av_packet_unref(packet);
		if (getFFMPEGLibraryVersion() >= 3)
		{
			av_packet_free(&packet);
		}
		else
		{
			av_free_packet(packet);
			packet = NULL;
		}
	}

	return frameLoaded;
}
void FileReader::close()
{
	if (vars->formatContext)
	{
		avformat_close_input(&vars->formatContext);
		vars->formatContext = NULL;
	}
}
Bool FileReader::peekFormat(Path filePath, FileFormat& format)
{
	FileReader peeker;
	peeker.vars->peekMode = TRUE;
	Bool ok = peeker.open(filePath, format);
	peeker.close();
	return ok;
}

FileWriter::FileWriter() : Object<FileWriterVars>(new FileWriterVars, TRUE)
{
	vars->formatContext = NULL;
	vars->writeHeaderLater = FALSE;
}
Bool FileWriter::open(Path filePath, FileFormat format)
{
	close();

	if (!filePath.isValid() || filePath.isFolder()) return FALSE;
	filePath.parentFolder().folderMake();

	String ext = filePath.extension().toLower();
	VideoContainerType containerType = VideoContainerType::Invalid;
	if (ext == ".avi") containerType = VideoContainerType::AVI;
	else if (ext == ".mp4") containerType = VideoContainerType::MP4;
	else if (ext == ".mkv") containerType = VideoContainerType::MKV;
	else return FALSE;

	AVCodecID codecID = AVCodecID::AV_CODEC_ID_NONE;
	if (format.codec == VideoDataCodec::H264) codecID = AVCodecID::AV_CODEC_ID_H264;
	else if (format.codec == VideoDataCodec::MJPEG) codecID = AVCodecID::AV_CODEC_ID_MJPEG;
	else if (format.codec == VideoDataCodec::YUV411) codecID = AVCodecID::AV_CODEC_ID_Y41P;
	else if (format.codec == VideoDataCodec::YUV420) codecID = AVCodecID::AV_CODEC_ID_YUV4;
	else if (format.codec == VideoDataCodec::H265) codecID = AVCodecID::AV_CODEC_ID_H265;
	else return FALSE;

	if (format.fps == 0) format.fps = 30;

	Pointer formatContext = NULL;
	avformat_alloc_output_context2(&formatContext, NULL, NULL, pathToUtf8CharArray(filePath).data());
	if (!formatContext) return FALSE;

	if (getFFMPEGLibraryVersion() >= 3)
	{
		Pointer stream = avformat_new_stream(formatContext, NULL);
		if (!stream)
		{
			avformat_free_context(formatContext);
			return FALSE;
		}

		avformat_stream_set_codec_type(stream, AVMediaType::AVMEDIA_TYPE_VIDEO);
		avformat_stream_set_codec_id(stream, codecID);
		avformat_stream_set_bit_rate(stream, 0);
		avformat_stream_set_width(stream, format.width);
		avformat_stream_set_height(stream, format.height);
		avformat_stream_set_time_base(stream, 1, format.fps);
		if (containerType == VideoContainerType::MP4) avformat_stream_set_time_base(stream, 1000, format.fps * 1000);
		if (format.codec == VideoDataCodec::H265) avformat_stream_set_codec_tag(stream, 0x31766568); // hev1
	}
	else
	{
		Pointer codec = avcodec_find_encoder(codecID);
		if (!codec)
		{
			avformat_free_context(formatContext);
			return FALSE;
		}

		Pointer stream = avformat_new_stream(formatContext, codec);
		if (!stream)
		{
			avformat_free_context(formatContext);
			return FALSE;
		}

		Pointer codecCtx = avformat_stream_get_codec(stream);
		avcodec_set_bit_rate(codecCtx, 0);
		avcodec_set_width(codecCtx, format.width);
		avcodec_set_height(codecCtx, format.height);
		avformat_stream_set_time_base(stream, 1, format.fps);
		if (containerType == VideoContainerType::MP4) avformat_stream_set_time_base(stream, 1000, format.fps * 1000);
		if (format.codec == VideoDataCodec::H265)
		{
			if (containerType == VideoContainerType::MP4) avcodec_set_codec_tag(codecCtx, 0x23); // hev1
			else avcodec_set_codec_tag(codecCtx, 0x31766568); // hev1
		}
		if (containerType == VideoContainerType::MP4 || containerType == VideoContainerType::MKV)
		{
			avcodec_set_flags(codecCtx, avcodec_get_flags(codecCtx) | AV_CODEC_FLAG_GLOBAL_HEADER);
		}
	}

	if (containerType == VideoContainerType::MKV)
	{
		if (format.codec == VideoDataCodec::H264) vars->writeHeaderLater = TRUE;
		else if (getFFMPEGLibraryVersion() >= 4 && format.codec == VideoDataCodec::H265) vars->writeHeaderLater = TRUE;
	}

	#if FFMPEGDL_VER_MAJOR == 1 && FFMPEGDL_VER_MINOR < 8
	if (vars->writeHeaderLater)
	{
		avformat_free_context(formatContext);
		return FALSE;
	}
	#endif

	Int resultOpen = avio_open(avformat_get_pb_ptr(formatContext), pathToUtf8CharArray(filePath).data(), AVIO_FLAG_WRITE);
	if (resultOpen < 0)
	{
		avformat_free_context(formatContext);
		return FALSE;
	}

	if (!vars->writeHeaderLater)
	{
		Int resultWriteHeader = avformat_write_header(formatContext, NULL);
		if (resultWriteHeader < 0)
		{
			avio_close(avformat_get_pb(formatContext));
			avformat_free_context(formatContext);
			return FALSE;
		}
	}

	vars->format = format;
	vars->containerType = containerType;
	vars->formatContext = formatContext;
	vars->count = 0;
	return TRUE;
}
UInt FileWriter::writeFrame(Binary frame)
{
	if (!vars->formatContext) return UINF;
	if (frame.size() < 8) return UINF;

	#if FFMPEGDL_VER_MAJOR > 1 || FFMPEGDL_VER_MINOR >= 8
	if (vars->writeHeaderLater)
	{
		Int extraDataSize = -1;
		if (vars->format.codec == VideoDataCodec::H264) extraDataSize = getH264ExtradataSize(frame);
		else if (vars->format.codec == VideoDataCodec::H265) extraDataSize = getH265ExtradataSize(frame);
		if (extraDataSize <= 0) return UINF;

		Pointer extraData = av_malloc(extraDataSize + AV_INPUT_BUFFER_PADDING_SIZE);
		utility::memorySet(0, extraData, extraDataSize + AV_INPUT_BUFFER_PADDING_SIZE);
		utility::memoryCopy(frame.data(), extraData, extraDataSize);

		Pointer stream = avformat_get_streams(vars->formatContext)[0];
		if (getFFMPEGLibraryVersion() >= 3)
		{
			avformat_stream_set_extradata(stream, extraData);
			avformat_stream_set_extradata_size(stream, extraDataSize);
		}
		else
		{
			Pointer codecCtx = avformat_stream_get_codec(stream);
			avcodec_set_extradata(codecCtx, extraData);
			avcodec_set_extradata_size(codecCtx, extraDataSize);
		}

		vars->writeHeaderLater = FALSE;

		Int resultWriteHeader = avformat_write_header(vars->formatContext, NULL);
		if (resultWriteHeader < 0)
		{
			avio_close(avformat_get_pb(vars->formatContext));
			avformat_free_context(vars->formatContext);
			vars->formatContext = NULL;
			return UINF;
		}
	}
	#endif

	Pointer packet = NULL;
	if (getFFMPEGLibraryVersion() >= 3)
	{
		packet = av_packet_alloc();
	}
	else
	{
		packet = av_mallocz_av_packet();
		av_init_packet(packet);
	}

	av_packet_set_dts(packet, vars->count);
	av_packet_set_pts(packet, vars->count);
	if (vars->containerType == VideoContainerType::MP4)
	{
		av_packet_set_dts(packet, (Long)vars->count * 1000);
		av_packet_set_pts(packet, (Long)vars->count * 1000);
	}
	else if (vars->containerType == VideoContainerType::MKV)
	{
		av_packet_set_dts(packet, (Long)(vars->count + 1) * 1000 / vars->format.fps);
		av_packet_set_pts(packet, (Long)(vars->count + 1) * 1000 / vars->format.fps);
	}

	av_packet_set_size(packet, frame.size());
	av_packet_set_data(packet, frame.data());

	int resultWriteFrame = -1;
	if (vars->format.codec == VideoDataCodec::H264 || vars->format.codec == VideoDataCodec::H265)
	{
		if ((vars->format.codec == VideoDataCodec::H264 && Frame::isDependentFrameH264(frame).isValidAndFalse()) ||
			(vars->format.codec == VideoDataCodec::H265 && Frame::isDependentFrameH265(frame).isValidAndFalse()))
		{
			av_packet_set_flags(packet, av_packet_get_flags(packet) | AV_PKT_FLAG_KEY);
		}
		resultWriteFrame = av_interleaved_write_frame(vars->formatContext, packet);
	}
	else
	{
		resultWriteFrame = av_write_frame(vars->formatContext, packet);
		vars->lastFrame = frame;
	}

	av_packet_unref(packet);
	if (getFFMPEGLibraryVersion() >= 3)
	{
		av_packet_free(&packet);
	}
	else
	{
		av_free_packet(packet);
		packet = NULL;
	}

	if (resultWriteFrame < 0) return UINF;
	else return vars->count++;
}
void FileWriter::close()
{
	if (vars->formatContext)
	{
		if (getFFMPEGLibraryVersion() >= 3 && !vars->lastFrame.isEmpty())
		{
			if (vars->format.codec == VideoDataCodec::MJPEG && vars->containerType == VideoContainerType::MP4)
			{
				writeFrame(vars->lastFrame);
			}
		}

		if (!vars->writeHeaderLater)
		{
			av_write_trailer(vars->formatContext);
		}

		avio_close(avformat_get_pb(vars->formatContext));
		avformat_free_context(vars->formatContext);
		vars->formatContext = NULL;
	}

	vars->writeHeaderLater = FALSE;
}